from pymilvus import MilvusClient, DataType, FieldSchema, CollectionSchema
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_milvus import Milvus
import pandas as pd
import streamlit as st

milvus_uri = st.secrets["milvus_uri"]
milvus_dbName = st.secrets["milvus_dbName"]
bailian_key = st.secrets["DASHSCOPE_API_KEY"]
collection_name = "interview_questions"

# 初始化 Milvus 客户端
client = MilvusClient(uri=milvus_uri, db_name=milvus_dbName)

# 定义 Schema（必须包含 vector 字段！）
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="type", dtype=DataType.VARCHAR, max_length=200),
    FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=1536)  # 维度需与 DashScopeEmbeddings 匹配
]
schema = CollectionSchema(fields)

# 初始化 Embedding 模型
embeddings = DashScopeEmbeddings(
    model="text-embedding-v2",
    max_retries=3,
    dashscope_api_key=bailian_key
)

# 初始化 LangChain 的 Milvus 向量库
vector_store = Milvus(
    embeddings,
    connection_args={"uri": milvus_uri,"db_name": milvus_dbName},
    collection_name=collection_name,
    text_field="content",      # 指定文本字段
    auto_id=True
)


def query(question):
    print(f"question:{question}")
    results = vector_store.similarity_search(question, k=2)
    print(f"vector_store查询结果:{results}")

def loadData(uploaded_file):
    df = pd.read_csv(uploaded_file)
    texts = df["content"].tolist()
    metadatas = [{"type": f"{uploaded_file.name}-{row['type']}"} for _, row in df.iterrows()]
    print(texts)
    print(metadatas)
    # 自动转向量并存储
    res = vector_store.add_texts(texts=texts, metadatas=metadatas)
    print(res)
    vector_store.col.flush()

    print(f"当前实体数: {vector_store.col.num_entities}")

def createCollection():
    index_params = client.prepare_index_params()
    index_params.add_index(
        field_name="embedding",
        index_type="IVF_FLAT",    # 量化索引，平衡速度与精度
        metric_type="L2",         # 相似性度量标准（欧式距离）
        params={"nlist": 1024}    # 聚类中心数
    )

    client.create_collection(
        collection_name=collection_name,
        schema=schema,
        index_params = index_params
    )